{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "ff223689",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "PreTrainedTokenizerFast(name_or_path='distilgpt2', vocab_size=50257, model_max_len=1024, is_fast=True, padding_side='right', truncation_side='right', special_tokens={'bos_token': '<|endoftext|>', 'eos_token': '<|endoftext|>', 'unk_token': '<|endoftext|>'})\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'input_ids': [[24717, 649, 3200, 507, 422, 262, 21694, 4991], [3642, 1299, 645, 20868, 837, 691, 2248, 1850, 308, 3775]], 'attention_mask': [[1, 1, 1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]}"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from transformers import AutoTokenizer\n",
    "\n",
    "#加载编码器\n",
    "tokenizer = AutoTokenizer.from_pretrained('distilgpt2', use_fast=True)\n",
    "\n",
    "print(tokenizer)\n",
    "\n",
    "#编码试算\n",
    "tokenizer.batch_encode_plus([\n",
    "    'hide new secretions from the parental units',\n",
    "    'contains no wit , only labored gags'\n",
    "])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c8a11a50",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at datas/glue/sst2/train/cache-2fa3d9ea0c2e3200.arrow\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at datas/glue/sst2/train/cache-cac77a5e86fd7743.arrow\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at datas/glue/sst2/train/cache-3cc26b7057143977.arrow\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at datas/glue/sst2/train/cache-8be3df80a5174f21.arrow\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at datas/glue/sst2/validation/cache-5defa0b47ec6fdf7.arrow\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at datas/glue/sst2/validation/cache-d1c27fdf594552d9.arrow\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at datas/glue/sst2/validation/cache-2839a3031c96535e.arrow\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at datas/glue/sst2/validation/cache-ca5072f02a45f80a.arrow\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at datas/glue/sst2/test/cache-328dd077154e4063.arrow\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at datas/glue/sst2/test/cache-787f5ae26460ac0c.arrow\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at datas/glue/sst2/test/cache-51866934df9b415a.arrow\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at datas/glue/sst2/test/cache-3b4181ace988b360.arrow\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at datas/glue/sst2/train/cache-2e9bc7b049aefd22_00000_of_00004.arrow\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at datas/glue/sst2/train/cache-2e9bc7b049aefd22_00001_of_00004.arrow\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at datas/glue/sst2/train/cache-2e9bc7b049aefd22_00002_of_00004.arrow\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at datas/glue/sst2/train/cache-2e9bc7b049aefd22_00003_of_00004.arrow\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at datas/glue/sst2/validation/cache-caae842e2bf7259e_00000_of_00004.arrow\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at datas/glue/sst2/validation/cache-caae842e2bf7259e_00001_of_00004.arrow\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at datas/glue/sst2/validation/cache-caae842e2bf7259e_00002_of_00004.arrow\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at datas/glue/sst2/validation/cache-caae842e2bf7259e_00003_of_00004.arrow\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at datas/glue/sst2/test/cache-2d531b689012f0db_00000_of_00004.arrow\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at datas/glue/sst2/test/cache-2d531b689012f0db_00001_of_00004.arrow\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at datas/glue/sst2/test/cache-2d531b689012f0db_00002_of_00004.arrow\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at datas/glue/sst2/test/cache-2d531b689012f0db_00003_of_00004.arrow\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at datas/glue/sst2/train/cache-e85f57c072cb0237.arrow\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at datas/glue/sst2/train/cache-c9f5532705e72c39.arrow\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at datas/glue/sst2/train/cache-d8ea0f09410378c4.arrow\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at datas/glue/sst2/train/cache-18ea0b109a1be497.arrow\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at datas/glue/sst2/validation/cache-98c60bb3a1be8cc2.arrow\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at datas/glue/sst2/validation/cache-5d48903f9f4f1a35.arrow\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at datas/glue/sst2/validation/cache-2e436811a1690c1e.arrow\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at datas/glue/sst2/validation/cache-beec3e430e9a2179.arrow\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at datas/glue/sst2/test/cache-e4318eba63d6359f.arrow\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at datas/glue/sst2/test/cache-37a41e2a435574e7.arrow\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at datas/glue/sst2/test/cache-74c7fa3e87d7b806.arrow\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at datas/glue/sst2/test/cache-8a1840d9a979acde.arrow\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(DatasetDict({\n",
       "     train: Dataset({\n",
       "         features: ['input_ids', 'attention_mask', 'labels'],\n",
       "         num_rows: 39905\n",
       "     })\n",
       "     validation: Dataset({\n",
       "         features: ['input_ids', 'attention_mask', 'labels'],\n",
       "         num_rows: 848\n",
       "     })\n",
       "     test: Dataset({\n",
       "         features: ['input_ids', 'attention_mask', 'labels'],\n",
       "         num_rows: 1730\n",
       "     })\n",
       " }),\n",
       " {'input_ids': [24717, 649, 3200, 507, 422, 262, 21694, 4991],\n",
       "  'attention_mask': [1, 1, 1, 1, 1, 1, 1, 1],\n",
       "  'labels': [24717, 649, 3200, 507, 422, 262, 21694, 4991]})"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from datasets import load_dataset, load_from_disk\n",
    "\n",
    "#加载数据\n",
    "#dataset = load_dataset(path='glue', name='sst2')\n",
    "dataset = load_from_disk('datas/glue/sst2')\n",
    "\n",
    "\n",
    "#分词,同时删除多余的字段\n",
    "def f(data):\n",
    "    return tokenizer.batch_encode_plus(data['sentence'])\n",
    "\n",
    "\n",
    "dataset = dataset.map(f,\n",
    "                      batched=True,\n",
    "                      batch_size=1000,\n",
    "                      num_proc=4,\n",
    "                      remove_columns=['sentence', 'idx', 'label'])\n",
    "\n",
    "\n",
    "#过滤掉太短的句子\n",
    "def f(data):\n",
    "    return [len(i) >= 8 for i in data['input_ids']]\n",
    "\n",
    "\n",
    "dataset = dataset.filter(f, batched=True, batch_size=1000, num_proc=4)\n",
    "\n",
    "\n",
    "#截断句子,同时整理成模型需要的格式\n",
    "def f(data):\n",
    "    data['input_ids'] = [i[:8] for i in data['input_ids']]\n",
    "    data['attention_mask'] = [[1] * 8] * len(data['attention_mask'])\n",
    "    #在模型中处理了偏移量问题,这里保持输入输出一致即可\n",
    "    data['labels'] = data['input_ids'].copy()\n",
    "    return data\n",
    "\n",
    "\n",
    "dataset = dataset.map(f, batched=True, batch_size=1000, num_proc=4)\n",
    "\n",
    "dataset, dataset['train'][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "405095b6",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(4988,\n",
       " {'input_ids': tensor([[10134,  1714,   319,  3159,   587,   523, 19855,  3098],\n",
       "          [33587,  2879,   318,   262,  1266,   339,   705,    82],\n",
       "          [ 2339,   427,  7274,   705,    82, 16924,   286,   262],\n",
       "          [10134,  1716,   262,  4958,   286,  3527,    84, 31110],\n",
       "          [  505,  1263, 12226,   284,   711,   530, 45179,  3715],\n",
       "          [ 1136,   517,   302,    12,  7513,   602,   286,   477],\n",
       "          [ 1462,   804, 45630,   272, 48973,   287,   262,  4151],\n",
       "          [  986,   743,  1234,   572, 31594,   290, 29028, 12936]]),\n",
       "  'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1, 1],\n",
       "          [1, 1, 1, 1, 1, 1, 1, 1],\n",
       "          [1, 1, 1, 1, 1, 1, 1, 1],\n",
       "          [1, 1, 1, 1, 1, 1, 1, 1],\n",
       "          [1, 1, 1, 1, 1, 1, 1, 1],\n",
       "          [1, 1, 1, 1, 1, 1, 1, 1],\n",
       "          [1, 1, 1, 1, 1, 1, 1, 1],\n",
       "          [1, 1, 1, 1, 1, 1, 1, 1]]),\n",
       "  'labels': tensor([[10134,  1714,   319,  3159,   587,   523, 19855,  3098],\n",
       "          [33587,  2879,   318,   262,  1266,   339,   705,    82],\n",
       "          [ 2339,   427,  7274,   705,    82, 16924,   286,   262],\n",
       "          [10134,  1716,   262,  4958,   286,  3527,    84, 31110],\n",
       "          [  505,  1263, 12226,   284,   711,   530, 45179,  3715],\n",
       "          [ 1136,   517,   302,    12,  7513,   602,   286,   477],\n",
       "          [ 1462,   804, 45630,   272, 48973,   287,   262,  4151],\n",
       "          [  986,   743,  1234,   572, 31594,   290, 29028, 12936]])})"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "from transformers.data.data_collator import default_data_collator\n",
    "\n",
    "#数据加载器\n",
    "loader = torch.utils.data.DataLoader(\n",
    "    dataset=dataset['train'],\n",
    "    batch_size=8,\n",
    "    collate_fn=default_data_collator,\n",
    "    shuffle=True,\n",
    "    drop_last=True,\n",
    ")\n",
    "\n",
    "for i, data in enumerate(loader):\n",
    "    break\n",
    "\n",
    "len(loader), data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "f53f3c45",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at distilgpt2 were not used when initializing GPT2Model: ['lm_head.weight']\n",
      "- This IS expected if you are initializing GPT2Model from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing GPT2Model from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "12050.9952\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(tensor(6.4238, grad_fn=<NllLossBackward0>), torch.Size([8, 8, 50257]))"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from transformers import AutoModelForCausalLM, GPT2Model\n",
    "\n",
    "#加载模型\n",
    "#model = AutoModelForCausalLM.from_pretrained('distilgpt2')\n",
    "\n",
    "\n",
    "#定义下游任务模型\n",
    "class Model(torch.nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.pretrained = GPT2Model.from_pretrained('distilgpt2')\n",
    "        self.fc = torch.nn.Linear(768, tokenizer.vocab_size, bias=False)\n",
    "        \n",
    "        #加载预训练模型的参数\n",
    "        parameters = AutoModelForCausalLM.from_pretrained('distilgpt2')\n",
    "        self.fc.load_state_dict(parameters.lm_head.state_dict())\n",
    "\n",
    "        self.criterion = torch.nn.CrossEntropyLoss()\n",
    "\n",
    "    def forward(self, input_ids, attention_mask, labels=None):\n",
    "        logits = self.pretrained(input_ids=input_ids,\n",
    "                                 attention_mask=attention_mask)\n",
    "        logits = logits.last_hidden_state\n",
    "\n",
    "        logits = self.fc(logits)\n",
    "\n",
    "        loss = None\n",
    "        if labels is not None:\n",
    "            shift_logits = logits[:, :-1].reshape(-1, tokenizer.vocab_size)\n",
    "            shift_labels = labels[:, 1:].reshape(-1)\n",
    "\n",
    "            loss = self.criterion(shift_logits, shift_labels)\n",
    "\n",
    "        return {'loss': loss, 'logits': logits}\n",
    "\n",
    "\n",
    "model = Model()\n",
    "\n",
    "#统计参数量\n",
    "print(sum(i.numel() for i in model.parameters()) / 10000)\n",
    "\n",
    "out = model(**data)\n",
    "\n",
    "out['loss'], out['logits'].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "6574726e",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "tensor([ 3159, 22152,  5321, 26810,   306,  1997,   318,   837])\n",
      "tensor([  995,  1943, 13839,  2646,   306,   257,    13,    11])\n",
      "10\n",
      "tensor([  803,  2004, 16304,   996,   463,   973,    12,   257])\n",
      "tensor([  803,    81,  2836,   290, 49191,   262,    12,   262])\n",
      "20\n",
      "tensor([39628,   287,   764,  7464,  1111,   356, 15827,  7357])\n",
      "tensor([ 267,  287,   13, 1110,  290,  407,  290, 7357])\n",
      "30\n",
      "tensor([ 2646,   262,  3227,   262, 23372,   290,    82,   663])\n",
      "tensor([ 13, 262,  13, 290, 265, 284,  82, 262])\n",
      "40\n",
      "tensor([  764,  1838,   475, 23365,   278,  2230,   422,   379])\n",
      "tensor([   13,   318,   290, 23365,   278,   837,   422,   290])\n",
      "50\n",
      "tensor([ 351, 3732,  389,  318,  290, 3613,  581,  262])\n",
      "tensor([  290, 42418,   284,   318,    11,  5368,   922,   262])\n",
      "0.20833333333333334\n",
      "plays as hollow catharsis,\n",
      " with  and\n",
      "\n",
      "somewhat blurred, but k\n",
      "inn udos\n",
      "\n",
      "denzel washington's efforts\n",
      " are  to\n",
      "\n",
      "a tale of horror and revenge that\n",
      " is  is\n",
      "\n",
      "these spiders can outrun a motorcycle\n",
      " and ,\n",
      "\n",
      "not even the hanson brothers can\n",
      " save  afford\n",
      "\n",
      "few films this year have been as\n",
      " res  good\n",
      "\n",
      "obvious politics and rudimentary animation reduce\n",
      " the  the\n",
      "\n"
     ]
    }
   ],
   "source": [
    "#测试\n",
    "def test():\n",
    "    model.eval()\n",
    "\n",
    "    #数据加载器\n",
    "    loader_test = torch.utils.data.DataLoader(\n",
    "        dataset=dataset['test'],\n",
    "        batch_size=8,\n",
    "        collate_fn=default_data_collator,\n",
    "        shuffle=True,\n",
    "        drop_last=True,\n",
    "    )\n",
    "\n",
    "    correct = 0\n",
    "    total = 0\n",
    "    for i, data in enumerate(loader_test):\n",
    "\n",
    "        #只计算最后一个词的正确率,这里先把最后一个词取出来\n",
    "        label = data['input_ids'][:, -1].clone()\n",
    "\n",
    "        #从数据中抹除掉最后一个词,防止模型作弊\n",
    "        data['input_ids'][:, -1] = 0\n",
    "\n",
    "        #label就不需要了\n",
    "        data['labels'][:, :] = 0\n",
    "\n",
    "        #计算\n",
    "        with torch.no_grad():\n",
    "            out = model(**data)\n",
    "\n",
    "        #只计算最后一个词的正确率,因为有偏移量的关系,这里取的是倒数第二个词\n",
    "        out = out['logits'].argmax(dim=2)[:, -2]\n",
    "\n",
    "        correct += (label == out).sum().item()\n",
    "        total += 8\n",
    "\n",
    "        if i % 10 == 0:\n",
    "            print(i)\n",
    "            print(label)\n",
    "            print(out)\n",
    "\n",
    "        if i == 50:\n",
    "            break\n",
    "\n",
    "    print(correct / total)\n",
    "\n",
    "    for i in range(8):\n",
    "        print(tokenizer.decode(data['input_ids'][i, :-1]))\n",
    "        print(tokenizer.decode(label[i]), tokenizer.decode(out[i]))\n",
    "        print()\n",
    "\n",
    "\n",
    "test()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "fb0e6ed9",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/root/anaconda3/envs/cpu/lib/python3.6/site-packages/transformers/optimization.py:309: FutureWarning: This implementation of AdamW is deprecated and will be removed in a future version. Use the PyTorch implementation torch.optim.AdamW instead, or set `no_deprecation_warning=True` to disable this warning\n",
      "  FutureWarning,\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 6.279179096221924 0.10714285714285714 1.999599037690457e-05\n",
      "50 5.47908353805542 0.21428571428571427 1.979550922213312e-05\n",
      "100 5.5193705558776855 0.16071428571428573 1.959502806736167e-05\n",
      "150 4.97205114364624 0.23214285714285715 1.939454691259022e-05\n",
      "200 5.653367042541504 0.19642857142857142 1.9194065757818766e-05\n",
      "250 4.709580421447754 0.25 1.8993584603047316e-05\n",
      "300 5.8663482666015625 0.16071428571428573 1.8793103448275863e-05\n",
      "350 4.809518337249756 0.19642857142857142 1.8592622293504413e-05\n",
      "400 4.790334224700928 0.26785714285714285 1.839214113873296e-05\n",
      "450 5.267386436462402 0.23214285714285715 1.819165998396151e-05\n",
      "500 4.75763463973999 0.26785714285714285 1.7991178829190057e-05\n",
      "550 5.370026111602783 0.17857142857142858 1.7790697674418608e-05\n",
      "600 4.932478427886963 0.2857142857142857 1.7590216519647154e-05\n",
      "650 4.6484599113464355 0.26785714285714285 1.7389735364875705e-05\n",
      "700 5.083922863006592 0.19642857142857142 1.718925421010425e-05\n",
      "750 5.430644989013672 0.16071428571428573 1.6988773055332802e-05\n",
      "800 5.452700614929199 0.16071428571428573 1.678829190056135e-05\n",
      "850 4.209372520446777 0.32142857142857145 1.6587810745789896e-05\n",
      "900 5.444892406463623 0.17857142857142858 1.6387329591018446e-05\n",
      "950 5.170382499694824 0.17857142857142858 1.6186848436246993e-05\n",
      "1000 4.150233745574951 0.30357142857142855 1.5986367281475543e-05\n",
      "1050 4.683593273162842 0.2857142857142857 1.578588612670409e-05\n",
      "1100 4.8343610763549805 0.23214285714285715 1.558540497193264e-05\n",
      "1150 6.111698150634766 0.08928571428571429 1.5384923817161187e-05\n",
      "1200 5.196045875549316 0.19642857142857142 1.5184442662389736e-05\n",
      "1250 4.255314350128174 0.3392857142857143 1.4983961507618286e-05\n",
      "1300 4.3259077072143555 0.2857142857142857 1.4783480352846833e-05\n",
      "1350 5.106658935546875 0.16071428571428573 1.4582999198075383e-05\n",
      "1400 4.975592613220215 0.23214285714285715 1.438251804330393e-05\n",
      "1450 4.998856067657471 0.32142857142857145 1.418203688853248e-05\n",
      "1500 4.919321537017822 0.23214285714285715 1.3981555733761027e-05\n",
      "1550 4.560444355010986 0.2857142857142857 1.3781074578989576e-05\n",
      "1600 4.621440410614014 0.21428571428571427 1.3580593424218124e-05\n",
      "1650 4.940247058868408 0.25 1.3380112269446673e-05\n",
      "1700 5.4352874755859375 0.14285714285714285 1.3179631114675222e-05\n",
      "1750 4.550337791442871 0.30357142857142855 1.297914995990377e-05\n",
      "1800 4.890089511871338 0.26785714285714285 1.2778668805132319e-05\n",
      "1850 4.5952372550964355 0.30357142857142855 1.2578187650360867e-05\n",
      "1900 5.1304097175598145 0.17857142857142858 1.2377706495589416e-05\n",
      "1950 4.011012077331543 0.30357142857142855 1.2177225340817965e-05\n",
      "2000 5.469344615936279 0.19642857142857142 1.1976744186046511e-05\n",
      "2050 4.604699611663818 0.26785714285714285 1.1776263031275062e-05\n",
      "2100 4.644625186920166 0.17857142857142858 1.1575781876503609e-05\n",
      "2150 4.466975212097168 0.32142857142857145 1.1375300721732159e-05\n",
      "2200 4.952640533447266 0.23214285714285715 1.1174819566960706e-05\n",
      "2250 4.5113091468811035 0.2857142857142857 1.0974338412189256e-05\n",
      "2300 5.101187229156494 0.19642857142857142 1.0773857257417803e-05\n",
      "2350 4.241959095001221 0.2857142857142857 1.0573376102646352e-05\n",
      "2400 4.415221691131592 0.26785714285714285 1.03728949478749e-05\n",
      "2450 5.049282550811768 0.125 1.0172413793103449e-05\n",
      "2500 4.529026031494141 0.2857142857142857 9.971932638331997e-06\n",
      "2550 3.9865596294403076 0.2857142857142857 9.771451483560546e-06\n",
      "2600 4.723902702331543 0.19642857142857142 9.570970328789094e-06\n",
      "2650 5.027300834655762 0.26785714285714285 9.370489174017643e-06\n",
      "2700 5.416659832000732 0.14285714285714285 9.170008019246192e-06\n",
      "2750 3.9759819507598877 0.4107142857142857 8.96952686447474e-06\n",
      "2800 5.1662492752075195 0.21428571428571427 8.769045709703289e-06\n",
      "2850 4.285780429840088 0.21428571428571427 8.568564554931837e-06\n",
      "2900 5.02206563949585 0.2857142857142857 8.368083400160386e-06\n",
      "2950 4.664642333984375 0.26785714285714285 8.167602245388933e-06\n",
      "3000 5.209836483001709 0.14285714285714285 7.967121090617481e-06\n",
      "3050 4.6904826164245605 0.21428571428571427 7.766639935846032e-06\n",
      "3100 4.853517055511475 0.23214285714285715 7.5661587810745795e-06\n",
      "3150 5.0579118728637695 0.19642857142857142 7.365677626303128e-06\n",
      "3200 3.844972848892212 0.375 7.165196471531677e-06\n",
      "3250 5.458155632019043 0.25 6.964715316760225e-06\n",
      "3300 5.093428134918213 0.19642857142857142 6.764234161988774e-06\n",
      "3350 3.9710516929626465 0.2857142857142857 6.563753007217322e-06\n",
      "3400 4.206329822540283 0.25 6.36327185244587e-06\n",
      "3450 4.449569225311279 0.25 6.162790697674419e-06\n",
      "3500 5.5621209144592285 0.19642857142857142 5.962309542902967e-06\n",
      "3550 4.712424278259277 0.2857142857142857 5.761828388131516e-06\n",
      "3600 4.3007307052612305 0.23214285714285715 5.5613472333600645e-06\n",
      "3650 4.432779312133789 0.19642857142857142 5.360866078588613e-06\n",
      "3700 3.8603622913360596 0.3392857142857143 5.1603849238171626e-06\n",
      "3750 4.192017555236816 0.32142857142857145 4.95990376904571e-06\n",
      "3800 4.681373119354248 0.23214285714285715 4.759422614274258e-06\n",
      "3850 4.424992084503174 0.25 4.558941459502807e-06\n",
      "3900 4.142448902130127 0.2857142857142857 4.358460304731355e-06\n",
      "3950 4.579164981842041 0.2857142857142857 4.157979149959905e-06\n",
      "4000 4.62972354888916 0.23214285714285715 3.957497995188452e-06\n",
      "4050 4.312018871307373 0.2857142857142857 3.757016840417001e-06\n",
      "4100 4.944700717926025 0.2857142857142857 3.5565356856455496e-06\n",
      "4150 4.435700416564941 0.30357142857142855 3.356054530874098e-06\n",
      "4200 4.214175701141357 0.25 3.1555733761026463e-06\n",
      "4250 4.892587184906006 0.14285714285714285 2.955092221331195e-06\n",
      "4300 4.3403120040893555 0.19642857142857142 2.7546110665597435e-06\n",
      "4350 4.7405829429626465 0.25 2.5541299117882925e-06\n",
      "4400 4.744002819061279 0.19642857142857142 2.3536487570168407e-06\n",
      "4450 4.451248645782471 0.25 2.153167602245389e-06\n",
      "4500 4.417544364929199 0.21428571428571427 1.952686447473938e-06\n",
      "4550 4.800900459289551 0.2857142857142857 1.752205292702486e-06\n",
      "4600 4.403017044067383 0.2857142857142857 1.5517241379310346e-06\n",
      "4650 5.2398481369018555 0.25 1.351242983159583e-06\n",
      "4700 4.581725120544434 0.3392857142857143 1.1507618283881318e-06\n",
      "4750 4.615292549133301 0.19642857142857142 9.502806736166801e-07\n",
      "4800 4.236312389373779 0.16071428571428573 7.497995188452287e-07\n",
      "4850 4.3059282302856445 0.23214285714285715 5.493183640737771e-07\n",
      "4900 5.0248236656188965 0.14285714285714285 3.488372093023256e-07\n",
      "4950 4.67105770111084 0.25 1.483560545308741e-07\n"
     ]
    }
   ],
   "source": [
    "from transformers import AdamW\n",
    "from transformers.optimization import get_scheduler\n",
    "\n",
    "\n",
    "#训练\n",
    "def train():\n",
    "    optimizer = AdamW(model.parameters(), lr=2e-5)\n",
    "    scheduler = get_scheduler(name='linear',\n",
    "                              num_warmup_steps=0,\n",
    "                              num_training_steps=len(loader),\n",
    "                              optimizer=optimizer)\n",
    "\n",
    "    model.train()\n",
    "    for i, data in enumerate(loader):\n",
    "        out = model(**data)\n",
    "        loss = out['loss']\n",
    "\n",
    "        loss.backward()\n",
    "        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n",
    "\n",
    "        optimizer.step()\n",
    "        scheduler.step()\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "        model.zero_grad()\n",
    "\n",
    "        if i % 50 == 0:\n",
    "            labels = data['labels'][:, 1:]\n",
    "            out = out['logits'].argmax(dim=2)[:, :-1]\n",
    "\n",
    "            correct = (labels == out).sum().item()\n",
    "            accuracy = correct / (8 * 7)\n",
    "\n",
    "            lr = optimizer.state_dict()['param_groups'][0]['lr']\n",
    "\n",
    "            print(i, loss.item(), accuracy, lr)\n",
    "\n",
    "    torch.save(model, 'models/1.预测最后一个词.model')\n",
    "\n",
    "\n",
    "train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "f8507bc2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "tensor([  268,   257,   837,   467, 31557,  3159,   326,  5730])\n",
      "tensor([ 268,  329,  290,  257, 3437, 1295,  286, 5730])\n",
      "10\n",
      "tensor([27235,   318,  1628, 17548,   257,   837,  1621,   783])\n",
      "tensor([ 837,  318,  837,  749,  705, 2003, 1621,  508])\n",
      "20\n",
      "tensor([2181,  318,  837,  768,  318,  286, 1628, 2085])\n",
      "tensor([3360,  318, 7328,  768,  318, 1022, 3807, 1178])\n",
      "30\n",
      "tensor([  416, 15360,   468,   837,   837,  7464,   705,   326])\n",
      "tensor([ 220,  257,  705,  220,  345, 3807,  705,  286])\n",
      "40\n",
      "tensor([ 290,  508, 1267,  290,  262,   12, 7675,  257])\n",
      "tensor([ 837, 2073, 1267,  220,  262,   12,  705,  257])\n",
      "50\n",
      "tensor([ 1372,   837,   502,   428,   546,   340, 16728,    82])\n",
      "tensor([262, 290, 262, 257, 326, 340, 286,  82])\n",
      "0.2818627450980392\n",
      "none of this violates the letter of\n",
      " beh  the\n",
      "\n",
      "chicago is sophisticated, brash\n",
      ",  and\n",
      "\n",
      "they felt like the same movie to\n",
      " me  the\n",
      "\n",
      "a recent favorite at sundance,\n",
      " this  a\n",
      "\n",
      "` dragonfly'is a movie\n",
      " about  that\n",
      "\n",
      "the problem with this film is that\n",
      " it  it\n",
      "\n",
      "if you're not the target\n",
      " demographic  of\n",
      "\n",
      "once the audience figure out what '\n",
      "s s\n",
      "\n"
     ]
    }
   ],
   "source": [
    "model = torch.load('models/1.预测最后一个词.model')\n",
    "test()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
